from PIL import Image
import numpy as np
import os
from torch.utils.data import Dataset
import math
import cv2
import torchvision.transforms as transforms
import torch
from data.base_dataset import BaseDataset, get_transform

class ColorAugmentation(object):
    def __init__(self):
        self.eig_vec = torch.Tensor([
            [0.4009, 0.7192, -0.5675],
            [-0.8140, -0.0045, -0.5808],
            [0.4203, -0.6948, -0.5836],
        ])
        self.eig_val = torch.Tensor([[0.2175, 0.0188, 0.0045]])

    def __call__(self, tensor):
        assert tensor.size(0) == 3
        alpha = torch.normal(mean=torch.zeros_like(self.eig_val)) * 0.1
        quatity = torch.mm(self.eig_val * alpha, self.eig_vec)
        tensor = tensor + quatity.view(3, 1, 1)
        return tensor

class CASIASURFdataset(BaseDataset):
    def __init__(self,opt):#traindir=None,valdir=None, transform=None, phase_train=True):
        self.npoints = None
        self.phase_train = (opt.controller == 'train')
        btoA = self.opt.direction == 'BtoA'
        input_nc = self.opt.output_nc if btoA else self.opt.input_nc       # get the number of channels of input image
        output_nc = self.opt.input_nc if btoA else self.opt.output_nc
        traindir = os.path.join(opt.dataroot,"Mytest_train.txt")
        valdir = os.path.join(opt.dataroot,"Mytest_test.txt")
        self.transform = get_transform(self.opt, grayscale=(input_nc == 1))
        
        self.channel = None
        try:
            if traindir is not None:
                with open(traindir, 'r') as f:
                    self.dir_train = f.read().splitlines()
                    self.rgb_dir_train = []
                    self.depth_dir_train = []
                    self.ir_dir_train = []
                    self.label_dir_train = []
                    for line in self.dir_train:
                        self.rgb_dir_train.append(line.split(',')[0])
                        self.depth_dir_train.append(line.split(',')[1])
                        self.ir_dir_train.append(line.split(',')[2])
                        self.label_dir_train.append(int(line.split(',')[3]))
                # with open(label_dir_train_file, 'r') as f:
                #     self.label_dir_train = f.read().splitlines()
            if valdir is not None:
                with open(valdir, 'r') as f:
                    self.dir_val = f.read().splitlines()
                    self.rgb_dir_test = []
                    self.depth_dir_test = []
                    self.ir_dir_test = []
                    self.label_dir_test = []
                    for line in self.dir_val:
                        self.rgb_dir_test.append(line.split(',')[0])
                        self.depth_dir_test.append(line.split(',')[1])
                        self.ir_dir_test.append(line.split(',')[2])
                        self.label_dir_test.append(int(line.split(',')[3]))
            self.cat_name = {}
                # with open(label_dir_val_file, 'r') as f:
                #     self.label_val = f.read().splitlines()
            # if self.phase_test:
            #     with open(depth_dir_test_file, 'r') as f:
            #         self.depth_dir_test = f.read().splitlines()
            #     with open(label_dir_test_file, 'r') as f:
            #         self.label_dir_test = f.read().splitlines()
        except:
            print('can not open files, may be filelist is not exist')
            exit()
    def set(self,phase_train=True):
        self.phase_train = phase_train
        
    def __len__(self):
        if self.phase_train:
            return len(self.dir_train)
        else:
            return len(self.dir_val)

    def __getitem__(self, idx):
        if self.phase_train:
            rgb_dir = self.rgb_dir_train[idx]
            depth_dir = self.depth_dir_train[idx]
            # label_dir = self.label_dir_train[idx]
            label = self.label_dir_train[idx]
            label = np.array(label)
        else:
            rgb_dir = self.rgb_dir_test[idx]
            depth_dir = self.depth_dir_test[idx]
            label = self.label_dir_test[idx]
            # label = int(label_dir[idx])
            label = np.array(label)

        rgb = Image.open(rgb_dir)
        rgb = rgb.convert('RGB')
        
        depth = Image.open(depth_dir)
        depth = depth.convert('RGB')

        # if self.transform:
        rgb = self.transform(rgb)
        depth = self.transform(depth)
        
        return {'S0':rgb,'depth':depth,'label':label,'isTrain':self.phase_train}
